package com.example.solrdemo.dao;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.solr.client.solrj.SolrRequest;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.request.RequestWriter;
import org.apache.solr.client.solrj.request.UpdateRequest;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.JavaBinCodec;
import org.apache.solr.common.util.NamedList;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.convert.converter.Converter;
import org.springframework.core.convert.converter.ConverterFactory;
import org.springframework.core.convert.support.ConfigurableConversionService;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.data.solr.core.SolrOperations;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Iterator;
import java.util.Map;

import static org.apache.solr.common.params.CommonParams.JAVABIN_MIME;

@Configuration
public class SolrConfig {

    public SolrConfig(SolrOperations operations, HttpSolrClient client) {

        client.setRequestWriter(jsonRequestWriter());

        ConfigurableConversionService service = (ConfigurableConversionService) operations.getConverter().getConversionService();

        service.addConverterFactory(new ConverterFactory<Object, Object>() {

            DefaultConversionService defaultConversionService = new DefaultConversionService();

            @Override
            public <T> Converter<Object, T> getConverter(Class<T> targetType) {

                if (defaultConversionService.canConvert(String.class, targetType)) {
                    return source -> defaultConversionService.convert(source, targetType);
                } else {
                    return source -> {
                        try {
                            return objectMapper().readValue((String) source, targetType);
                        } catch (IOException e) {
                            return null;
                        }
                    };
                }

            }
        });
    }

    @Bean
    public ObjectMapper objectMapper() {
        return new ObjectMapper();
    }

    @Bean
    public JsonRequestWriter jsonRequestWriter() {
        return new JsonRequestWriter(objectMapper());
    }

    static class JsonRequestWriter extends RequestWriter {

        private MyJavaBinUpdateRequestCodec codec;

        public JsonRequestWriter(ObjectMapper objectMapper) {
            codec = new MyJavaBinUpdateRequestCodec(objectMapper);
        }

        @Override
        public ContentWriter getContentWriter(SolrRequest req) {

            if (req instanceof UpdateRequest) {
                UpdateRequest updateRequest = (UpdateRequest) req;
                if (isEmpty(updateRequest)) return null;
                return new ContentWriter() {
                    @Override
                    public void write(OutputStream os) throws IOException {
                        codec.marshal(updateRequest, os);
                    }

                    @Override
                    public String getContentType() {
                        return JAVABIN_MIME;
                    }
                };
            } else {
                return req.getContentWriter(JAVABIN_MIME);
            }
        }


    }

    static class MyJavaBinUpdateRequestCodec {


        private ObjectMapper objectMapper;

        private final JavaBinCodec.ObjectResolver objectResolver = (o, codec) -> objectMapper.writeValueAsString(o);

        public MyJavaBinUpdateRequestCodec(ObjectMapper objectMapper) {
            this.objectMapper = objectMapper;
        }

        void marshal(UpdateRequest updateRequest, OutputStream os) throws IOException {
            NamedList<Object> nl = new NamedList<>();
            NamedList<Object> params = solrParamsToNamedList(updateRequest.getParams());
            if (updateRequest.getCommitWithin() != -1) {
                params.add("commitWithin", updateRequest.getCommitWithin());
            }
            Iterator<SolrInputDocument> docIter = null;

            if (updateRequest.getDocIterator() != null) {
                docIter = updateRequest.getDocIterator();
            }

            Map<SolrInputDocument, Map<String, Object>> docMap = updateRequest.getDocumentsMap();

            nl.add("params", params);// 0: params
            if (updateRequest.getDeleteByIdMap() != null) {
                nl.add("delByIdMap", updateRequest.getDeleteByIdMap());
            }
            nl.add("delByQ", updateRequest.getDeleteQuery());

            if (docMap != null) {
                nl.add("docsMap", docMap.entrySet().iterator());
            } else {
                if (updateRequest.getDocuments() != null) {
                    docIter = updateRequest.getDocuments().iterator();
                }
                nl.add("docs", docIter);
            }
            try (JavaBinCodec codec = new JavaBinCodec(objectResolver)) {
                codec.marshal(nl, os);
            }
        }


        private NamedList<Object> solrParamsToNamedList(SolrParams params) {
            if (params == null) return new NamedList<>();
            return params.toNamedList();
        }

    }


}
